package trace

import (
	"context"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// OTELInterceptorBuilder 用于构建 OpenTelemetry 拦截器。
type OTELInterceptorBuilder struct {
	tracer      trace.Tracer
	propagator  propagation.TextMapPropagator
	serviceName string
}

// NewOTELInterceptorBuilder 创建一个 OTELInterceptorBuilder 实例。
func NewOTELInterceptorBuilder(
	serviceName string,
	tracer trace.Tracer,
	propagator propagation.TextMapPropagator) *OTELInterceptorBuilder {
	return &OTELInterceptorBuilder{tracer: tracer,
		serviceName: serviceName, propagator: propagator}
}

// BuildUnaryServerInterceptor 服务端的拦截器
func (b *OTELInterceptorBuilder) BuildUnaryServerInterceptor() grpc.UnaryServerInterceptor {
	propagator := b.propagator
	if propagator == nil {
		propagator = otel.GetTextMapPropagator()
	}
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (reply interface{}, err error) {
		ctx = extract(ctx, propagator)
		return handler(ctx, req)
	}
}

// BuildUnaryClientInterceptor 客户端的拦截器。
func (b *OTELInterceptorBuilder) BuildUnaryClientInterceptor() grpc.UnaryClientInterceptor {
	propagator := b.propagator
	if propagator == nil {
		propagator = otel.GetTextMapPropagator()
	}
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) {
		ctx = inject(ctx, propagator)
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}

// extract 从 gRPC 请求的上下文中提取传播的追踪信息。
func extract(ctx context.Context, propagators propagation.TextMapPropagator) context.Context {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = metadata.MD{}
	}
	return propagators.Extract(ctx, GrpcHeaderCarrier(md))
}

// inject 将传播的追踪信息注入到 gRPC 请求的上下文中。
func inject(ctx context.Context, propagators propagation.TextMapPropagator) context.Context {
	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		md = metadata.MD{}
	}
	propagators.Inject(ctx, GrpcHeaderCarrier(md))
	return metadata.NewOutgoingContext(ctx, md)
}

// GrpcHeaderCarrier 是 OpenTelemetry 的传播接口的 gRPC 头部载体实现。
type GrpcHeaderCarrier metadata.MD

// Get 返回与传递的键关联的值。
func (mc GrpcHeaderCarrier) Get(key string) string {
	vals := metadata.MD(mc).Get(key)
	if len(vals) > 0 {
		return vals[0]
	}
	return ""
}

// Set 存储键值对。
func (mc GrpcHeaderCarrier) Set(key string, value string) {
	metadata.MD(mc).Set(key, value)
}

// Keys 列出此载体中存储的键。
func (mc GrpcHeaderCarrier) Keys() []string {
	keys := make([]string, 0, len(mc))
	for k := range metadata.MD(mc) {
		keys = append(keys, k)
	}
	return keys
}
